# -*- coding:utf-8 -*-
"""
    公共模块
"""
import sys,os
from pprint import pprint
import os, sys, json, re, time, pwd, logging,logging.handlers
from textwrap import dedent

try:
    from urllib.request   import urlopen
    from urllib.request   import Request
    from urllib.parse     import quote,urlencode,urlparse
except:
    from urllib2    import urlopen
    from urllib2    import Request
    from urllib     import quote,urlencode
    from urlparse   import urlparse

# from urllib import parse
# from urllib import request

def color(string, code=2):
    return "\033[1;3{}m{}\033[0m".format(code, string)

def green(string):
    return color(string, 2)

def red(string):
    return color(string, 1)

def yellow(string):
    return color(string, 3)

def init_log(filename="/tmp/anscan.log"):
    filename = filename
    logger = logging.getLogger(__name__)
    # 日志输出到控制台
    stdout = logging.StreamHandler()
    stdout.setFormatter(logging.Formatter('[%(asctime)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S'))
    logger.addHandler(stdout)
    # fh = logging.FileHandler(filename, mode='w+', encoding='utf-8')
    # logger.addHandler(fh)
    logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(message)s", datefmt='%Y-%m-%d %H:%M:%S', filename=filename, filemode='a')
    return logger

logger = init_log()

def log(*info):
    logger.info(" ".join([str(_) for _ in info]))

def anscan_filelog_path(option):
    if not os.path.isdir(option.dirname):
        print("------------------------------")
        print("请正确的输入目录：")
        return
    else:
        logpath = os.path.join(option.dirname, "anscan.log")
        g_anscan_filepath["anscanlog"] = logpath

# report.txt 报告输出
def init_log_report(filename="/tmp/report.txt"):
    loggerr = logging.getLogger(filename)
    # loggerr.addHandler()
    # logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(message)s", datefmt='%Y-%m-%d %H:%M:%S', filename=filename, filemode='a')


    fh = logging.FileHandler(filename, mode='w+')
    ch = logging.StreamHandler()
    formatter = logging.Formatter('[%(asctime)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
    ch.setFormatter(formatter)
    # fh.setFormatter(formatter)
    loggerr.addHandler(fh)
    loggerr.addHandler(ch)
    return loggerr




def delete(fname):
    if not os.path.exists(fname):
        return True
    return os.unlink(fname)

def tar(filelist=[]):
    tmp = tempfile.mktemp()
    open(tmp, 'w+').write("\n".join(filelist))
    output = "{}.tar.gz".format(int(time.time()))
    output = os.path.join(os.getcwd(), output)
    cmd = "tar -zcvf {} --files-from {} 2>/dev/null".format(output, tmp)
    ret = os.popen(cmd).read()
    # log(ret)
    # delete(tmp)
    return output

# 计算样本MD5
def calculation_md5(filelist):
    output = os.path.join("/tmp/", "filemd5.txt")
    # cmd = '" " > {}'.format(output)
    # ret = os.popen(cmd)
    open(output, "w").close()
    for key in filelist:
        if os.path.exists(key):
            try:
                cmd = "md5sum {} 2>/dev/null >> {}".format(key, output)
                ret = os.popen(cmd).read()
            except Exception as e:
                log("md5sum 命令对改文件权限不足")
                cmd = "该文件{} 因权限不足无法计算MD5值  >> {}".format(key, output)
                os.popen(cmd).read()
    return output
# 记录设备信息
def device_info():
    cmd = "netstat -antulp > /tmp/netstat && ps auxf > /tmp/ps && cat /proc/cpuinfo > /tmp/cpu && cat /proc/meminfo > /tmp/mem && lsblk > /tmp/disk"
    ret = os.popen(cmd).read()
    return ret

def config(enable=True, alias="test", **options):
    import functools
    def _config(function):
        @functools.wraps(function)
        def wrapper(option, *args, **kwargs):
            if option.no and wrapper.alias in option.no:
                return None
            if option.only and wrapper.alias not in option.only:
                return None
            log("----------------------------------------------")
            log("【开始执行】{}".format(function.__doc__))
            try:
                ret = function(option, *args, **kwargs)
            except:
                # 检测失败
                log(dumpstack())
                ret = False
            if ret:
                log(red("【有风险】 {}".format(function.__doc__)))
            else:
                log(green("【通过】 {}".format(function.__doc__)))
            log("------------------end------------------------")
            return ret
        setattr(wrapper, "enable", enable)
        setattr(wrapper, "alias", alias.lower())
        for k,v in options.items():
            setattr(wrapper, k, v)
        return wrapper
    return _config

def ls(_dir):
    return [os.path.join(_dir, _) for _ in os.listdir(_dir)]

def cat(fname):
    if not os.path.exists(fname): return ""
    return open(fname).read()

# 取文件前10M内容
def head(file, max_size="10M"):
    if not os.path.exists(file): return ""
    if os.path.isdir(file): return ""
    if (" " in file) or ("\\" in file) or (".jpg" in file) or (")" in file) or ("(" in file) or (".log" in file): return ""
    if (os.path.getsize(file) == 0): return ""
    strings = os.popen("head -c %s '%s' 2>/dev/null | strings 2>/dev/null" % (max_size, file)).read().splitlines()
    return strings

# 不做处理，记录威胁的文件名
def danger_file(*files):
    for flink in files:
        file = os.path.realpath(flink)
        if file not in g_dangerous_files:
            log(red("发现可疑文件 {}".format(file)))
            g_dangerous_files[file] = 1
            continue
        g_dangerous_files[file] += 1

def monitoring_objects_num(obj):
    g_monitoring_objects_num["count"] += 1
    if obj in g_monitoring_objects_num:
        g_monitoring_objects_num[obj] += 1
    else:
        g_monitoring_objects_num[obj] = 1

def save_result(desc, info, level="high"):
    if level.lower()[0] == 'h':
        # 高风险
        log(red("[高风险]"), desc, info)
        g_message_high.append((desc, info))
    elif level.lower()[0] == 'm':
        # 中风险
        log(red("[中风险]"), desc, info)
        g_message_medium.append((desc, info))
    elif level.lower()[0] == 'l':
        # 低风险
        log(yellow("[低风险]"),desc, info)
        g_message_low.append((desc, info))

# 分析字符串是否包含境外IP  存在境外IP匹配返回真  不存在境外ip返回假
def check_contents_ip(contents):
    # try:
    if 1:
        if not re.search(ip_http, contents): return False
        if re.search(lan_ip, contents): return False
        for ip in re.findall(ip_re, contents):
            if (find(ip)[0:2] != u'中国') and (find(ip)[0:3] != u'局域网') and (find(ip)[0:4] != u'共享地址') and (
                    find(ip)[0:4] != u'本机地址') and (find(ip)[0:4] != u'本地链路') and (find(ip)[0:4] != u'保留地址'):
                return True
        return False
    # except:
    #     return False

# 判断是否为ip 是ip 返回真 非ip 返回假
def isIP(str):
    p = re.compile('^((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)$')
    if p.match(str):
        return True
    else:
        return False

# 检测IP是否境外IP 是境外ip则返回真  否则返回假
def check_ip(ip):
    # try:
    if 1:
        ip = ip.strip()
        # query_args = {"ip": "69.174.99.181"}
        query_args = {"ip": ip}
        encode_args = parse.urlencode(query_args).encode('utf-8')
        url = 'http://api.e4ting.cn/v2/api/engine'
        result = request.urlopen(url, encode_args).read().decode('utf-8')
        result = json.loads(result)
        if result["data"]["info"][0]["report"]["malname"]:
            return True
        else:
            return False

# 递归目录返回文件名列表
def gci(filepath):
    filename = []
    try:
    # if 1:
        files = os.listdir(filepath)
        for fi in files:
            fi_d = os.path.join(filepath, fi)
            if os.path.isdir(fi_d):
                filename = filename + gci(fi_d)
            else:
                filename.append(os.path.join(filepath, fi_d))
        return filename
    except:
        return filename

# 分析一串字符串是否包含反弹shell、获取对应字串内可能存在的文件，并判断文件是否存在恶意特征。
# 匹配成功则返回恶意特征信息
# 否则返回空
def analysis_strings(contents):
    try:
    # if 1:
        content = contents.replace('\n', '')
        # 反弹shell类
        if check_shell(content):
            return u"反弹shell类：%s" % content
        # 境外IP操作类
        #elif check_contents_ip(content):
        #    return u"境外ip操作类：%s" % content
        else:
            for file in content.split(' '):
                if type(file) == "byte": continue
                file = file.split(";")[0]
                if not os.path.exists(file):
                    continue
                if os.path.isdir(file):
                    continue
                malware = analysis_file(file)
                if malware: return u"引用恶意文件%s，可疑内容：%s" % (file, malware)
        return ""
    except:
        return ""

# 分析字符串是否包含反弹shell或者恶意下载执行的特征
def check_shell(content):
    # try:
    if 1:
        # 反弹shell类
        if (
            ('bash' in content) and
            (
                ('/dev/tcp/' in content) or
                ('telnet ' in content) or
                ('nc ' in content) or
                (('exec ' in content) and ('socket' in content)) or
                ('curl ' in content) or
                ('wget ' in content) or
                ('lynx ' in content) or
                ('bash -i' in content)
            )
           ) or (".decode('base64')" in content) or ("exec(base64.b64decode" in content):
            return content

        elif ('/dev/tcp/' in content) and (('exec ' in content) or ('ksh -c' in content)):
            return content
        elif ('exec ' in content) and (('socket.' in content) or (".decode('base64')" in content)):
            return content
        # 下载执行类
        elif (('wget ' in content) or ('curl ' in content)) and (
                (' -O ' in content) or (' -s ' in content)
            ) and (' http' in content) and (
                ('php ' in content) or
                ('perl' in content) or
                ('python ' in content) or
                ('sh ' in content) or
                ('bash ' in content)
             ):
            return content
        return False
    # except:
    #     return False

# 分析文件是否包含恶意特征、反弹shell特征、境外ip类信息
# 存在返回恶意特征
# 不存在返回空
def analysis_file(file, mode='fast'):
    # try:
    if 1:
        strings =  head(file)
        if len(strings) > 200: return ""

        time.sleep(0.01)
        for str in strings:
            if check_shell(str):
                log("%s 文件存在反弹shell后门"%(file))
                danger_file(file) # shell 反弹有结果之后之直接记录到g_dangerous_files = {}
                # log(u'文件：%s ，bash shell :%s' % (file, str))
                # return u"反弹shell"
                return True
        return ""
    # except:
    #     log(dumpstack())
    #     return ""

# 检测配置文件是否存在恶意配置
def check_conf(tag, file, mode='only'):
    # print("check_conf", file, 'export ' + tag)
    # try:
    if 1:
        if not os.path.exists(file): return ""
        if os.path.isdir(file): return ""
        monitoring_objects_num(tag)
        if mode == 'only':
            with open(file) as f:
                for line in f:
                    line = line.strip()
                    # print("export", line)
                    if len(line) < 3: continue
                    if line[0] == '#': continue
                    if 'export ' + tag in line:
                        return line
        else:
            return analysis_file(file)
        return ""
    # except:
    #     return ""

# 检测所有环境变量
def check_tag(name, tag, mode='only'):
    suspicious = False
    # try:
    if 1:
        files = ['/root/.bashrc', '/root/.tcshrc', '/root/.bash_profile', '/root/.cshrc', '/root/.tcshrc',
                 '/etc/bashrc', '/etc/profile', '/etc/profile.d/', '/etc/csh.login', '/etc/csh.cshrc']
        home_files = ['/.bashrc', '/.bash_profile', '/.tcshrc', '/.cshrc', '/.tcshrc']

        # 循环用户目录查看环境设置
        for dir in os.listdir('/home/'):
            for home_file in home_files:
                file = os.path.join('%s%s%s' % ('/home/', dir, home_file))
                info = check_conf(tag, file, mode)
                if info == True:
                    save_result(name, "%s 设置了 后门代码" % (file), 'h')
                    suspicious = True
                elif info:
                    danger_file(file)
                    danger_file(*info.split("=")[-1].split(";"))
                    save_result(name, "%s 设置了 %s " % (file, info), 'm')
                    suspicious = True
        # 检查系统目录的配置
        for file in files:
            # 如果为目录形式，则遍历目录下所有文件
            if os.path.isdir(file):
                for _file in ls(file):
                    info = check_conf(tag, _file, mode)
                    if info == True:
                        save_result(name, "%s 设置了 后门代码" % (file), 'h')
                        suspicious = True
                    elif info:
                        danger_file(_file)
                        danger_file(*info.split("=")[-1].split(";"))
                        # save_result(u'常规后门检测', name, _file, '', info, u'[1]echo $%s [2]cat %s' % (tag, _file),u'可疑')
                        save_result(name, "%s 设置了 %s " % (_file, info), 'm')
                        suspicious = True
            else:
                info = check_conf(tag, file, mode)
                if info == True:
                    save_result(name, "%s 设置了 后门代码" % (file), 'h')
                    suspicious = True
                elif info:
                    danger_file(file)
                    danger_file(*info.split("=")[-1].split(";"))
                    save_result(name, "%s 设置了 %s " % (file, info), 'm')
                    suspicious = True
        return suspicious
    # except:
    #     return suspicious

# correct_baopo_infos = []
# 循环遍历/var/log 目录获取 sevure 类似的文件
def dir_file_detect(log_dir):
    correct_baopo_infos = []
    files = [os.path.join(log_dir, i) for i in os.listdir(log_dir) if
                (not os.path.isdir(i)) and ('secure' in i)]
    for log in files:
        monitoring_objects_num("sshbreak") # 每个文件都加一个检测项，这里考虑是否只需要添加一个检测项即可
        attack_detect(log, correct_baopo_infos)
    return correct_baopo_infos
# 数组去重
def reRepeat(old):
    new_li = []
    for i in old:
        if i not in new_li:
            new_li.append(i)
    return new_li

# 过滤
def _filter(old, count):
    new_li = []
    for key in old:
        if old[key] > count:
            new_li.append({key: old[key]})
    return new_li

# 实现counter函数，统计ip出现的次数，返回一个dict
def Counter(old):
    count_dict = dict()
    for item in old:
        if item in count_dict:
            count_dict[item] += 1
        else:
            count_dict[item] = 1
    return count_dict

def attack_detect(log, correct_baopo_infos):
    # 账户错误特征
    username_error = 'Invalid user'

    # correct_baopo_infos = correct_baopo_infos
    # 账户正确密码错误特征
    username_correct = 'Failed password for'
    # 成功登陆
    username_password_correct = 'Accepted password for'
    # 所有错误登陆日志ip
    failed_ip = []
    # 登陆成功日志
    correct_infos = []
    # C段ip登陆错误日志
    failed_c_ips = []
    filename = os.path.basename(log)
    year = ''
    if 'secure-' in filename and len(filename) == 15:
        year = filename[7:11]
    # 打开日志文件
    try:
        f = open(log, 'r')
    except Exception as e:
        return
    for i in f:
        if (username_error in i) and ('from' in i) and ('sshd' in i):
            # try:
            if 1:
                failed_ip.append(i.split(': ')[1].split()[4])
            # except:
            #     continue
        elif (username_correct in i) and ('from' in i) and ('sshd' in i):
            # try:
            if 1:
                failed_ip.append(i.split(': ')[1].rsplit()[-4])
            # except:
            #     continue
        elif username_password_correct in i and ('sshd' in i):
            ip = i.split(': ')[1].split()[5]
            user = i.split(': ')[1].split()[3]
            # time = i.split(' sshd[')[0]
            time = ' '.join(i.replace('  ', ' ').split(' ', 4)[:3]) + " " + year
            # 获取所有登陆成功的记录
            correct_infos.append({'ip': ip, 'user': user, 'time': time})
    # 记录登陆失败攻击源IP地址和尝试次数
    # 1.1 判断是否发生了爆破行为,failed_ip_dict为存在爆破的失败ip列表:次数
    failed_ip_dict = _filter(dict(Counter(failed_ip)), ip_failed_count)
    # 1.2 判断是否发生了C段类的爆破行为，

    for key in failed_ip:
        failed_c_ips.append(key.rsplit('.', 1)[0])
    failed_c_ips_dict = _filter(dict(Counter(failed_c_ips)), ips_failed_count)

    # 2、判断爆破行为是否成功，爆破成功才能算数
    for correct_info in correct_infos:
        for failed in failed_ip_dict:
            if correct_info['ip'] in failed: correct_baopo_infos.append(correct_info)
        for failed in failed_c_ips_dict:
            if correct_info['ip'].rsplit('.', 1)[0] in failed: correct_baopo_infos.append(correct_info)
    correct_baopo_infos = reRepeat(correct_baopo_infos)

def calc_rate(level="l", n=0):
    if n <= 0:
        return 0
    from math import e
    if level == "l":
        ret = round(100/( (1 + e **(55 - 0.75*n) )**0.1 ) , 1)
    elif level == "m":
        ret = round(100/( (1 + e **(35 - 0.75*n) )**0.1 ) , 1)
    elif level == "h":
        ret = round(100/( (1 + e **(23 - 0.75*n) )**0.1 ) , 1)
    else:
        ret = 0
    return ret

